function [phi,psi,keff,Q] = cmfd_eigen(phi,psi,in,dx,fmmid,Q,fd,keff,mu,w,nfm,ncm)
% function [psi,phi] = cmfd_eigen(phi,psi,in,dx,fmmid,Q,mu,wt,nfm,ncm)
% This function performs one outer iteration in coarse mesh finite
% difference space following the formulation given by Zhong et al. in
% NSE(158).
%
%   Inputs:
%           psi     -- fine mesh angular flux
%           phi     -- fine mesh scalar flux
%           Q       -- fine mesh external source
%           in      -- inpute structure
%           mtt     -- material index for each fine mesh
%           mu      -- angle set
%           w       -- angle weight
%           dx      -- fine mesh delta x
%   Output
%           psi     -- updated fine mesh angular flux
%           phi     -- updated fine mesh scalar flux

% allocate group dependent data
N  = length(in.xfm);                  % number of coarse meshes
sT = zeros( N, in.numg );             % homogenized total cross-section
sF = zeros( N, in.numg );             % homogenized nuSigF cross-section
sX = zeros( N, in.numg );             % homogenized Chi
sS = zeros( N, in.numg, in.numg  );   % homogenized scattering cross-section (group-to-group)
%                       (...,g1) sSg1->g1  sSg1->g2 ...
%                        ...,g2) sSg2->g1  ....
D  = zeros( N, in.numg );             % homogenized diffusion coefficient
Dt = zeros( N+1, in.numg );           % "D tilde"
Dh = zeros( N+1, in.numg );           % "D hat" (current correction factor)
FD  = zeros( N, 1 );                  % region fission density
S  = zeros( N,in.numg );              % region source
J  = zeros( N+1, in.numg );           % net current at edges
phi_ref  = zeros( N, in.numg );       % reference coarse mesh scalar flux

% coarse mesh delta x's
h = in.xcm(2:end)-in.xcm(1:end-1);

% the angle-integrated fission source
s = zeros(sum(in.xfm),in.numg);  
for i = 1:length(s)
    for g = 1:in.numg
        s(i,g) = sum( Q(i,:,g)'.*w(:) );
    end
end

for g = 1:in.numg
    J(1) = sum( mu .* w.* psi(1,:,g)' ); % first edge net current
end
for i = 1:length(in.xfm)  
    % note, the homogenization is meaningless until I implement a way to
    % specify coarse meshes that contain heterogeneities
    idx1        = 1 + sum( in.xfm(1:(i-1)) );          % lower index
    idx2        = sum( in.xfm(1:(i  )) );              % upper index   
    
    for g = 1:in.numg
        J(i+1,g)      = sum( mu .* w.* psi(idx2+1,:,g)' );     % net current at edges
        phi_ref(i,g)  = sum( dx(idx1:idx2)'.*phi(idx1:idx2,g) ) / h(i);
        for gg = 1:in.numg
            sS(i,g,gg) = sum( dx(idx1:idx2)'.*phi(idx1:idx2,g)* ...
                in.data((in.mt(i)-1)*in.numg+gg,4+g))/(phi_ref(i)*h(i));
        end
        sT(i,g)       = sum( dx(idx1:idx2)'.*phi(idx1:idx2,g)*...
                        in.data( in.mt(i), 1 ))/(phi_ref(i)*h(i));        
        sF(i,g)       = sum( dx(idx1:idx2)'.*phi(idx1:idx2,g)*...
                        in.data( in.mt(i), 3 ))/(phi_ref(i)*h(i));
        sX(i,g)       = sum( dx(idx1:idx2)'.*phi(idx1:idx2,g)*...
                        in.data( in.mt(i), 4 ))/(phi_ref(i)*h(i));
        D(i,g)        = 1/(3*sT(i));
        S(i,g)        = sum( dx(idx1:idx2)'.*s(idx1:idx2,g) ) / h(i);    
    end
    FD(i,g)        = sum( dx(idx1:idx2)'.*fd(idx1:idx2,1) ) / h(i); % homogenized fission density    
end
a = 0.5;
% for g = 1:in.numg
%     Dt(1,g)   = 2*a*D(1,g) / (a*h(1)+2*D(1,g));
%     Dt(end,g) = -2*a*D(end,g) / (a*h(end)+2*D(end,g));
%     Dh(1,g)   = -(J(1,g)+Dt(1,g)*phi_ref(1,g))/phi_ref(1,g);
%     Dh(end,g) = -(J(end,g)+Dt(end,g)*phi_ref(end,g))/phi_ref(end,g);
%     for i = 2:length(Dh(:,1))-1
%        Dt(i,g) = 2*D(i-1,g)*D(i,g)/(D(i-1,g)*h(i)+D(i,g)*h(i-1));
%        Dh(i,g) = - ( J(i,g) + Dt(i,g)*(phi_ref(i,g)-phi_ref(i-1,g)) ) / (phi_ref(i,g)+phi_ref(i-1,g)) ;
%     end
% end

% our cmfd implentation is based on the expression
%   Jnet_{i+1/2) = -Dt_{i+1/2}*(phi_{i+1}-phi_{i}) 
%                  - Dh_{i+1/2}*(phi_{i+1}+phi_{i}) 
% where Dt is an effective diffusion coefficient (based on the 
% homegenized D_i and D_{i+1}) and Dh is the current correction
% coefficient.  Since we get the Jnet's from the fine mesh calculation,
% we define Dh so that the equation above is satisfied.
% 
Dt(1)   = 2*a*D(1) / (a*h(1)+2*D(1));         %  Dt(1)   = 1;
Dt(end) = -2*a*D(end) / (a*h(end)+2*D(end));  %  Dt(end) = 1;
Dh(1)   = -(J(1)+Dt(1)*phi_ref(1))/phi_ref(1);
Dh(end) = -(J(end)+Dt(end)*phi_ref(end))/phi_ref(end);
for i = 2:length(Dh)-1
   Dt(i) = 2*D(i-1)*D(i)/(D(i-1)*h(i)+D(i)*h(i-1));
  % Dt(i) = 1;
   Dh(i) = - ( J(i) + Dt(i)*(phi_ref(i)-phi_ref(i-1)) ) / (phi_ref(i)+phi_ref(i-1)) ;
end

% L C U (diagonals of matrix A, for A*phi = S)
C = zeros(N,in.numg );
L = zeros(N-1,in.numg );
U = L;
for g = 1:in.numg    
    % left edge
    C(1,g) = 1/h(1) * ( Dt(2,g) - Dh(2,g) + Dt(1,g) + Dh(1,g) ) + sT(1,g)-sS(1,g,g);
    U(1,g) = 1/h(1) * (-Dt(2,g)-Dh(2,g));
    % right edge
    C(end,g) = 1/h(end) * ( -Dt(end,g) + Dt(end-1,g) - Dh(end,g) + Dh(end-1,g) ) + sT(end,g)-sS(end,g,g);
    L(end,g) = 1/h(end) * (-Dt(end-1,g)+Dh(end-1,g));
    % internal
    for i = 2:length(C)-1
        C(i,g)   = 1/h(i) * (Dt(i+1,g)+Dt(i,g)-Dh(i+1,g)+Dh(i,g)) + sT(i,g)-sS(i,g,g);  % C * phi(i)
        U(i,g)   = 1/h(i) * (-Dt(i+1,g)-Dh(i+1,g));                                     % U * phi(i+1)
        L(i-1,g) = 1/h(i) * (-Dt(i,g)+Dh(i,g));                                         % L * phi(i-1)
    end
end

for i = 1:N
    m = in.mt(i);
    for g = 1:in.numg
        S(i,g) = in.data((m-1)*in.numg+g, 4 ) * FD(i);
    end
end
S_ref = S;
phi_c = phi_ref;
errK = 1; epsK = 1e-9; errS = 1; epsS = 1e-9; iter = 0; % careful with these hardcoded convergence criteria
% begin outers
while   ( (( errK>epsK ) && (errS>epsS)) && iter < 6000 )
    % SOLVE FLUXES
    for i = 1:N
        m = in.mt(i);
        for g = 1:in.numg
            S(i,g) = in.data((m-1)*in.numg+g, 4 ) * FD(i);
        end
    end    
    phi_c(:,1)  = tridiag(U(:,1),L(:,1),C(:,1),S(:,1)/keff);
    scsrc       = zeros(N,1);
    if in.numg > 1                  % successive groups
        for g = 2:in.numg
            for gg = 1:(g-1)        
                for i = 1:N         % compute scattering source
                    scsrc(i,1) = scsrc(i,1) + sS(i,gg,g)*phi_c(i,gg);
                end
            end 
            % solve for phi_g
            phi_c(:,g)    = tridiag(U(:,1),L(:,1),C(:,1),scsrc+S(:,g)/keff);
        end
    end
    % save old fission density
    FDo = FD;
    kold = keff;   
    FD(:) = 0;       % reset fission source
    for g = 1:in.numg  % compute fission source
      FD(:) = FD(:) + sF(:,g).*phi_c(:,g);
    end  
    % compute keff
    keff = sum(FD.*h')*kold/sum(FDo.*h');
    
    % convergence stuff
    errS = max( abs((FD-FDo)./FDo) );
    errK = abs( (keff-kold)/kold );
    iter = iter + 1;  % number of iterations
end
   % disp( ['    FINAL CMFD keff =',num2str(keff),' in ITS=',num2str(iter)] )
    sprintf('   FINAL CMFD keff = %2.8f in %i',keff,iter )
% update:  
%    within meshes    
%           psi(ii,m) = psi(ii,m)*f(i),    ii E i
%    add coarse mesh edges
%           psi(ii,m) = psi(ii,m)*f(i),    ii E S(i,i')  and mu > 0
%           psi(ii,m) = psi(ii,m)*f(i'),   ii E (Si,i')  and mu < 0
for i = 1:N  
    % within cell indices
    idx1 = 1 + sum( in.xfm(1:(i-1)) ); % lower index
    idx2 = sum( in.xfm(1:(i  )) );     % upper index   
    for g = 1:in.numg
        phi(idx1:idx2,g)   = phi(idx1:idx2,g)  * phi_c(i,g)/phi_ref(i,g);
    end
    Q(idx1:idx2,:,g) = Q(idx1:idx2,:,g) * S(i,g)/S_ref(i,g);
end

end